package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{VectorIndexEncoder, VectorIndexEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.linalg.Vectors
import org.scalatest.FunSuite

/**
  * @Author tangjinyuan
  * @Date 2019/3/27 14:32
  * @Description
  * @Reviewer
  */
class VectorIndexEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "features",
      |	"outputCol": "indexedFeatures",
      |	"maxCategories": 5,
      | "preserveCols":"features,indexedFeatures"
      |}
    """.stripMargin

  val model = new VectorIndexEncoder()
  val param: VectorIndexEncoderParam = model.parseParam(new ToolClass().encrypt(paramStr))

  test("model parameter") {
    assert("features".equals(param.inputCol) && "indexedFeatures".equals(param.outputCol) && param.maxCategories == 5)
  }


  val seq = Seq(
    (Vectors.dense(Array(2.0, 5.0, 7.0)).toSparse, ""),
    (Vectors.dense(Array(3.0, 5.0, 9.0)).toSparse, ""),
    (Vectors.dense(Array(4.0, 7.0, 9.0)).toSparse, ""),
    (Vectors.dense(Array(2.0, 4.0, 9.0)).toSparse, ""),
    (Vectors.dense(Array(9.0, 5.0, 7.0)).toSparse, ""),
    (Vectors.dense(Array(2.0, 5.0, 9.0)).toSparse, ""),
    (Vectors.dense(Array(3.0, 4.0, 9.0)).toSparse, ""),
    (Vectors.dense(Array(8.0, 4.0, 9.0)).toSparse, ""),
    (Vectors.dense(Array(3.0, 6.0, 2.0)).toSparse, ""),
    (Vectors.dense(Array(5.0, 9.0, 2.0)).toSparse, "")
  )

  val compareSeq = Seq(
    (Vectors.sparse(3, Array(0, 1, 2), Array(2.0, 1.0, 1.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(3.0, 1.0, 2.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(4.0, 3.0, 2.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(2.0, 0.0, 2.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(9.0, 1.0, 1.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(2.0, 1.0, 2.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(3.0, 0.0, 2.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(8.0, 0.0, 2.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(3.0, 2.0, 0.0)), ""),
    (Vectors.sparse(3, Array(0, 1, 2), Array(5.0, 4.0, 0.0)), "")
  )

  test("model transformation") {
    val data = TestingDataGenerator.spark.createDataFrame(seq).toDF("features", "b")
    val compareData = TestingDataGenerator.spark.createDataFrame(compareSeq).toDF("indexedFeatures", "b")
    println("compareData:")
    compareData.select("indexedFeatures").show()
    val res = model.train(param, data).transform(data)
    println("res:")
    res.select("indexedFeatures").show()


    assert(compareData.select("indexedFeatures").except(res.select("indexedFeatures")).count() == 0)
  }


}
